#!/usr/bin/env python3
"""
PyTorch CNN图像识别器
支持多种输入格式：原始图像、numpy数组、PyTorch tensor
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
from PIL import Image
import numpy as np
import os
from typing import Union, List


class CNNImageClassifier(nn.Module):
    """简单的CNN图像分类器"""
    
    def __init__(self, num_classes: int = 10, input_channels: int = 3):
        super(CNNImageClassifier, self).__init__()
        
        # 第一个卷积块
        self.conv1 = nn.Conv2d(input_channels, 32, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.pool1 = nn.MaxPool2d(2)
        
        # 第二个卷积块
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.conv4 = nn.Conv2d(128, 128, kernel_size=3, padding=1)
        self.bn4 = nn.BatchNorm2d(128)
        self.pool2 = nn.MaxPool2d(2)
        
        # 第三个卷积块
        self.conv5 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.bn5 = nn.BatchNorm2d(256)
        self.conv6 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.bn6 = nn.BatchNorm2d(256)
        self.pool3 = nn.MaxPool2d(2)
        
        # 全连接层
        self.dropout = nn.Dropout(0.5)
        self.fc1 = nn.Linear(256 * 4 * 4, 512)  # 假设输入为32x32，经过3次池化后为4x4
        self.fc2 = nn.Linear(512, num_classes)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播"""
        # 第一个卷积块
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = self.pool1(x)
        
        # 第二个卷积块
        x = F.relu(self.bn3(self.conv3(x)))
        x = F.relu(self.bn4(self.conv4(x)))
        x = self.pool2(x)
        
        # 第三个卷积块
        x = F.relu(self.bn5(self.conv5(x)))
        x = F.relu(self.bn6(self.conv6(x)))
        x = self.pool3(x)
        
        # 全连接层
        x = x.view(x.size(0), -1)  # 展平
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        
        return x


class ImageProcessor:
    """图像处理器，支持多种输入格式"""
    
    def __init__(self, image_size: tuple = (32, 32)):
        self.image_size = image_size
        self.transform = transforms.Compose([
            transforms.Resize(image_size),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        
    def process(self, input_data: Union[str, np.ndarray, torch.Tensor, Image.Image]) -> torch.Tensor:
        """
        处理多种格式的输入数据
        
        Args:
            input_data: 支持的文件路径、numpy数组、PyTorch tensor、PIL图像
            
        Returns:
            torch.Tensor: 处理后的图像tensor，形状为(1, 3, H, W)
        """
        
        # 1. 文件路径输入
        if isinstance(input_data, str):
            if os.path.isfile(input_data):
                image = Image.open(input_data).convert('RGB')
                return self.transform(image).unsqueeze(0)  # 添加batch维度
            else:
                raise FileNotFoundError(f"文件不存在: {input_data}")
        
        # 2. PIL图像输入
        elif isinstance(input_data, Image.Image):
            image = input_data.convert('RGB')
            return self.transform(image).unsqueeze(0)
        
        # 3. numpy数组输入
        elif isinstance(input_data, np.ndarray):
            # 检查数组形状和数据类型
            if input_data.dtype != np.uint8:
                input_data = (input_data * 255).astype(np.uint8)
            
            # 处理不同的数组形状
            if len(input_data.shape) == 2:  # 灰度图 (H, W)
                input_data = np.stack([input_data] * 3, axis=-1)  # 转为RGB
            elif len(input_data.shape) == 3:
                if input_data.shape[2] == 1:  # 单通道 (H, W, 1)
                    input_data = np.concatenate([input_data] * 3, axis=2)
                elif input_data.shape[2] == 4:  # RGBA (H, W, 4)
                    input_data = input_data[:, :, :3]  # 转为RGB
            
            # 转换为PIL图像再处理
            image = Image.fromarray(input_data)
            return self.transform(image).unsqueeze(0)
        
        # 4. PyTorch tensor输入
        elif isinstance(input_data, torch.Tensor):
            # 确保tensor在正确的设备上且数据类型正确
            if input_data.device.type != 'cpu':
                input_data = input_data.cpu()
            
            # 处理不同的tensor形状
            if len(input_data.shape) == 2:  # (H, W)
                input_data = input_data.unsqueeze(0).repeat(3, 1, 1)  # (3, H, W)
            elif len(input_data.shape) == 3:
                if input_data.shape[0] == 1:  # (1, H, W)
                    input_data = input_data.repeat(3, 1, 1)  # (3, H, W)
                elif input_data.shape[0] == 4:  # (4, H, W)
                    input_data = input_data[:3]  # (3, H, W)
            
            # 添加batch维度并调整大小
            if len(input_data.shape) == 3:
                input_data = input_data.unsqueeze(0)  # (1, 3, H, W)
            
            # 调整大小到目标尺寸
            if input_data.shape[2:] != self.image_size:
                input_data = F.interpolate(input_data, size=self.image_size, mode='bilinear', align_corners=False)
            
            # 标准化
            mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
            std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
            input_data = (input_data - mean) / std
            
            return input_data
        
        else:
            raise TypeError(f"不支持的输入类型: {type(input_data)}")


class ImageRecognizer:
    """图像识别器，整合模型和处理器"""
    
    def __init__(self, num_classes: int = 10, image_size: tuple = (32, 32)):
        self.num_classes = num_classes
        self.image_size = image_size
        self.processor = ImageProcessor(image_size)
        self.model = CNNImageClassifier(num_classes)
        
        # 加载预训练权重（如果有）
        self._load_pretrained_weights()
        
        # 设置为评估模式
        self.model.eval()
        
    def _load_pretrained_weights(self):
        """加载预训练权重（这里为示例，实际使用时需要真实权重文件）"""
        # 这里可以加载预训练模型权重
        # 例如：self.model.load_state_dict(torch.load('model_weights.pth'))
        pass
        
    def recognize(self, input_data: Union[str, np.ndarray, torch.Tensor, Image.Image, List]) -> dict:
        """
        识别图像内容
        
        Args:
            input_data: 支持多种输入格式，也支持批量输入
            
        Returns:
            dict: 识别结果，包含预测类别和置信度
        """
        
        # 处理批量输入
        if isinstance(input_data, list):
            return self._recognize_batch(input_data)
        
        # 处理单个输入
        with torch.no_grad():
            # 预处理输入数据
            processed_tensor = self.processor.process(input_data)
            
            # 模型预测
            outputs = self.model(processed_tensor)
            probabilities = F.softmax(outputs, dim=1)
            confidence, predicted_class = torch.max(probabilities, 1)
            
            return {
                'class': predicted_class.item(),
                'confidence': confidence.item(),
                'all_probabilities': probabilities.squeeze().tolist()
            }
    
    def _recognize_batch(self, input_list: List) -> List[dict]:
        """批量识别"""
        results = []
        for input_data in input_list:
            results.append(self.recognize(input_data))
        return results